# -*- coding:utf-8 -*-
import os,sys
import copy
import re
import traceback
import time
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), os.pardir, os.pardir))
import supeanut_config
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), os.pardir))
from CommonLib.mylog import mylog
from BasicTool import BasicTool
from CommonLib.StockTool import StockTool
from CommonLib.db.MongoDBTool import MongoDBTool


'''
作者：supeanut
创建时间：2016-xx-xx xx:xx:xx
功能描述：
	xxx
	xxxxx
相关配置：
	supeanut_config.XXX
历史改动：
	2016-xx-xx: xxxxxx
'''
class ExrightsTool:
	# init what
	def __init__(self):
		self.stocktool = StockTool()
		self.basictool = BasicTool()
		self.mongoDB = supeanut_config.MONGO_DB
		self.exrights_coll = supeanut_config.MONGO_COLL_EXRIGHTS

	# input: code"300033"
	# output: {"2017-02-23":{送股，增股，派息，配股，配股价}}
	def getExrights(self, code):
		# 从mongo拿取除权数据
		mongodb_obj = MongoDBTool(None)
		flag, msg = mongodb_obj.conn(self.mongoDB, self.exrights_coll)
		if flag is False:
			return False, msg
		try:
			docs_iter = mongodb_obj.coll.find({"code":code})
		except:
			return False, traceback.format_exc()
		# 除复权数据str->float
		result_dict = {}
		for doc in docs_iter:
			date = doc["date"].encode("utf-8")
			songgu = float(doc["songgu"]) if doc["songgu"] != "--" else 0.0
			zenggu = float(doc["zenggu"]) if doc["zenggu"] != "--" else 0.0
			paixi = float(doc["paixi"]) if doc["paixi"] != "--" else 0.0
			peigu = float(doc["peigu"]) if doc["peigu"] != "--" else 0.0
			peigujia = float(doc["peigujia"]) if doc["peigujia"] != "--" else 0.0
			result_dict[date] = {"songgu":songgu, "zenggu":zenggu, "paixi":paixi, \
								 "peigu":peigu, "peigujia":peigujia}
		return True, result_dict

	# 若拿取局部时间区域内的数据，复权结果不为全局结果
	# 大盘无需除复权，因此code为6位简码，不包含指数代码
	def getAdjItemList(self, code, mod="pre", period="day", start_datetime=None, end_datetime=None):
		# 获取ori数据
		flag, oriItemList = self.basictool.getOriItemList(code, False, period, start_datetime, end_datetime)
		if flag is False:
			return False, oriItemList
		if len(oriItemList) == 0:
			return True, []
		# 获取exrights数据
		flag, exrights = self.getExrights(code)
		if flag is False:
			return False, exrights
		if exrights == {}:
			return True, oriItemList
		# 计算复权数据
		if mod == "pre":
			flag, AdjItemList = self.preCal(oriItemList, exrights)
		elif mod == "post":
			flag, AdjItemList = self.postCal(oriItemList, exrights)
		elif mod == "ori":
			return True, oriItemList
		else:
			return False, "mod param should be in ori,pre or post"
		if flag is False:
			return False, AdjItemList
		return True, AdjItemList

	def postCal(self, oriItemList, exrights):
		AdjItemList = copy.deepcopy(oriItemList)
		# 计算exrights的datelist
		all_exrights_datelist = list(exrights)
		exrights_datelist = []
		# --排除日期区间外的数据
		s_date = oriItemList[0][0][:10]
		e_date = oriItemList[-1][0][:10]
		for exrights_date in all_exrights_datelist:
			if exrights_date >= s_date and exrights_date <= e_date:
				exrights_datelist.append(exrights_date)
		if len(exrights_datelist) == 0:
			return True, oriItemList
		exrights_datelist.sort()
		exrights_datelist_len = len(exrights_datelist)
		# 计算复权累加区域
		period_dict = {}
		search_exright_date = exrights_datelist[0]
		period_index = 0
		period_index_max = 0
		period_end = len(oriItemList) - 1
		item_index = -1
		for item in oriItemList:
			item_index += 1
			# 去除datetime中的time
			date = item[0][:10]
			if date >= search_exright_date:
			#if date in exrights:
				period_dict[period_index] = {"period":[item_index, period_end], \
											 "exrights_date":search_exright_date}
				period_index_max = period_index
				period_index += 1
				if period_index >= exrights_datelist_len:
					break
				search_exright_date = exrights_datelist[period_index]
		# 如果start-end区间内无除复权数据，返回原itemList
		if period_dict == {}:
			return True, oriItemList
		# 倒序累加
		for period_index in range(period_index_max, -1, -1):
			period = period_dict[period_index]["period"]
			exrights_date = period_dict[period_index]["exrights_date"]
			exrights_data = exrights[exrights_date]
			paixi = exrights_data["paixi"]
			peigu = exrights_data["peigu"]
			peigujia = exrights_data["peigujia"]
			songgu = exrights_data["songgu"]
			zenggu = exrights_data["zenggu"]
			# 复权效果累加
			for item in AdjItemList[period[0]: period[1]+1]:
				# 价格复权
				for i in [1,2,3,4]:
					item[i] = item[i] * (1.0 + songgu + zenggu + peigu) - peigu * peigujia + paixi
				# 成交量复权
				item[5] = item[5] / (1.0 + songgu + zenggu + peigu)
		return True, AdjItemList
	
	def preCal(self, oriItemList, exrights):
		AdjItemList = copy.deepcopy(oriItemList)
		# 计算exrights的datelist
		all_exrights_datelist = list(exrights)
		exrights_datelist = []
		# --排除日期区间外的数据
		s_date = oriItemList[0][0][:10]
		e_date = oriItemList[-1][0][:10]
		for exrights_date in all_exrights_datelist:
			if exrights_date >= s_date and exrights_date <= e_date:
				exrights_datelist.append(exrights_date)
		if len(exrights_datelist) == 0:
			return True, oriItemList
		exrights_datelist.sort()
		exrights_datelist_len = len(exrights_datelist)
		# 计算复权累加区域
		period_dict = {}
		search_exright_date = exrights_datelist[0]
		period_index = 0
		period_index_max = 0
		period_start = 0
		item_index = -1
		for item in oriItemList:
			item_index += 1
			# 去除datetime中的time
			date = item[0][:10]
			if date >= search_exright_date:
			#if date in exrights:
				period_dict[period_index] = {"period":[period_start, item_index-1], \
											 "exrights_date":search_exright_date}
				period_index_max = period_index
				period_index += 1
				if period_index >= exrights_datelist_len:
					break
				search_exright_date = exrights_datelist[period_index]
		# 按历史复权顺序，进行复权累加
		#                historyN ...historyN-1... cur data
		# one fuquan: <------>
		# two fuquan: <------------------>
		# N fuquan:   <--------------------------->
		# 如果start-end区间内无除复权数据，返回原itemList
		if period_dict == {}:
			return True, oriItemList
		for period_index in range(0, period_index_max + 1):
			period = period_dict[period_index]["period"]
			exrights_date = period_dict[period_index]["exrights_date"]
			exrights_data = exrights[exrights_date]
			paixi = exrights_data["paixi"]
			peigu = exrights_data["peigu"]
			peigujia = exrights_data["peigujia"]
			songgu = exrights_data["songgu"]
			zenggu = exrights_data["zenggu"]
			# 复权效果累加
			for item in AdjItemList[period[0]: period[1]+1]:
				# 价格复权
				for i in [1,2,3,4]:
					item[i] = (item[i] - paixi + peigu * peigujia) / ( 1.0 + songgu + zenggu + peigu)
				# 成交量复权
				item[5] = item[5] * (1.0 + songgu + zenggu + peigu)
		return True, AdjItemList

if __name__ == '__main__':
	obj = ExrightsTool()
	print obj.getAdjItemList("000035", "pre", "day", "2016-01-01", "2017-03-28")
